import sys
import os

# 获取项目根目录的路径（即包含 scripts 和 modules 的目录）
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(project_root)
import lcm
import toml
import sys
import time
import numpy as np
from threading import Thread

import rclpy
import torch
from rclpy.qos import QoSProfile, QoSHistoryPolicy
from rclpy.node import Node
from sensor_msgs.msg import LaserScan, Image
from std_msgs.msg import Float64

from utils.state_estimator_lcmt import state_estimator_lcmt
from utils.robot_control_response_lcmt import robot_control_response_lcmt
from utils.leg_control_data_lcmt import leg_control_data_lcmt
from utils.motor_ctrl_lcmt import motor_ctrl_lcmt

import argparse
from utils.motor_robot_controller import Motor_Robot_Controller
from net import model
import torch.nn.functional as F
from obs_buffer import *

msg = motor_ctrl_lcmt()
Ctrl = Motor_Robot_Controller()
Ctrl.run()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
start_time = time.time()

def sendZeroFrames():
    for _ in range(8):
        msg.q_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        msg.qd_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        msg.kp_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        msg.kd_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        msg.tau_des = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        Ctrl.Send_cmd(msg)
        time.sleep(0.5)
        print("sending zero frame.")


class ObsSubscriber(Node):
    def __init__(self):
        super().__init__("obs_subscriber")

        self.handle_imu_thread = Thread(target=self.handle_imu)
        self.handle_contact_thread = Thread(target=self.handle_contact)
        self.handle_leg_thread = Thread(target=self.handle_leg)
        self.handle_send_cmd_thread = Thread(target=self.handle_send_cmd)
        self.lcm_imu = lcm.LCM("udpm://239.255.76.67:7669?ttl=255")
        self.lcm_contact = lcm.LCM("udpm://239.255.76.67:7670?ttl=255")
        self.lcm_leg = lcm.LCM("udpm://239.255.76.67:7667?ttl=255")

        self.lcm_imu.subscribe("state_estimator", self.handle_imu_msg)
        self.lcm_contact.subscribe("robot_control_response", self.handle_contact_state)
        self.lcm_leg.subscribe("leg_control_data", self.handle_leg_state)

        qos_profile = QoSProfile(
            reliability=rclpy.qos.QoSReliabilityPolicy.BEST_EFFORT,
            history=QoSHistoryPolicy.KEEP_LAST,
            depth=5,
        )
        self.depth_subscription = self.create_subscription(
            Image,
            "/camera/depth/image_rect_raw",
            self.depth_callback,
            qos_profile=qos_profile,
        )
        # self.subscription = self.create_subscription(
        #     Float64,
        #     '/rl_rate',
        #     self.rl_callback,
        #     10)
        # self.timer_ = self.create_timer(0.02, self.rl_callback)
        self.thread_get_data = Thread(target=self.rl_callback)
        # imu数据, 角速度, 加速度
        self.omega = []
        self.acc = []
        # 关节角度
        self.q = []
        # 关节角速度
        self.qd = []
        # 触地
        self.contact = []
        # 电机控制消息
        self.motor_msg = motor_ctrl_lcmt()
        # 网络模型
        self.model_use = model(
            args_use.n_proprio,
            args_use.n_scan,
            args_use.n_actions,
            args_use.n_priv_latent,
            args_use.n_priv_explicit,
            args_use.n_hist,
        )

        # obs的创建
        self.obs_class = ObservationBuffer(1, args_use.n_obs)

        # 深度图    具体怎么写我还没看   假设获得的是[1, 58, 87]
        self.depth_image = []
        print("Initialize done.")

    def handle_imu_msg(self, channel, data):
        # 50hz
        msg = state_estimator_lcmt.decode(data)
        self.omega = msg.omegaWorld
        self.acc = msg.aWorld

    def handle_contact_state(self, channel, data):
        # 50hz
        msg = robot_control_response_lcmt().decode(data)
        binary_str = f"{msg.contact:04b}"

        # contact 中的顺序: 左后, 右后, 左前, 右前
        self.contact = [0.5 if digit == "0" else -0.5 for digit in binary_str]

    def handle_leg_state(self, channel, data):
        msg = leg_control_data_lcmt().decode(data)
        self.q = msg.q
        self.qd = msg.qd

    def handle_imu(self):
        try:
            while True:
                self.lcm_imu.handle()
        except KeyboardInterrupt:
            print("Stopping subscriber.")

    def handle_contact(self):
        try:
            while True:
                self.lcm_contact.handle()
        except KeyboardInterrupt:
            print("Stopping subscriber.")

    def handle_leg(self):
        try:
            while True:
                self.lcm_leg.handle()
        except KeyboardInterrupt:
            print("Stopping subscriber.")

    def handle_send_cmd(self):
        try:
            while True:
                Ctrl.Send_cmd(self.motor_msg)
                self.q = Ctrl.currentQ
                self.qd = Ctrl.currentQd
        except KeyboardInterrupt:
            print("Stopping subscriber.")

    def run(self):
        self.handle_imu_thread.start()
        self.handle_contact_thread.start()
        self.handle_leg_thread.start()
        self.handle_send_cmd_thread.start()

        self.thread_get_data.start()

    # def get_need_data(self):
    #     # 我怎么感觉这个都不需要了
    #     while True:
    #         obs_data = []
    #         # 添加imu数据   6个
    #         if self.omega and self.acc:
    #             obs_data.extend(self.omega)
    #             obs_data.extend(self.acc)
    #         # 网络获得的数据  2个   3个0
    #
    #         # 关节角度和角速度
    #         if self.q and self.qd:
    #             obs_data.extend(self.q)
    #             obs_data.extend(self.qd)
    #
    #         # 12个动作
    #
    #         # 添加触地数据
    #         if self.contact:
    #             obs_data.extend(self.contact)
    #         obs_tensor = torch.tensor(obs_data, dtype=torch.float32).view(1, -1)
    #         self.obs_class.obs_buf[:, : args_use.n_proprio] = obs_tensor
    #         # # 添加imu数据   6个
    #         # if self.omega and self.acc:
    #         #     self.obs_class.obs_buf[:, : 3] = torch.tensor(self.omega, dtype=torch.float32).view(1, -1)
    #         #     self.obs_class.obs_buf[:, 3: 6] = torch.tensor(self.acc, dtype=torch.float32).view(1, -1)
    #         #
    #         # # 关节角度和角速度
    #         # if self.q and self.qd:
    #         #     self.obs_class.obs_buf[:, 13: 25] = torch.tensor(self.q, dtype=torch.float32).view(1, -1)
    #         #     self.obs_class.obs_buf[:, 25: 37] = torch.tensor(self.qd, dtype=torch.float32).view(1, -1)
    #         #
    #         # # 添加触地数据
    #         # if self.contact:
    #         #     self.obs_class.obs_buf[:, 49: 53] = torch.tensor(self.contact, dtype=torch.float32).view(1, -1)

    # 获得深度图是以10hz的速率运行
    def depth_callback(self, msg):
        print(time.time()-start_time)
        data = np.frombuffer(msg.data, dtype=np.uint16)
        image = np.reshape(data, (msg.height, msg.width))
        tensor_image = (
            torch.tensor(image, dtype=torch.float32, device=device).unsqueeze(0).unsqueeze(0)
        )
        downsampled_image = F.interpolate(
            tensor_image, size=(58, 87), mode="bilinear", align_corners=False
        )

        self.depth_image = downsampled_image.view(1, 58, 87)
        self.obs_class.depth_latent_and_yaw = self.model_use.depth_model(self.depth_image)
        print(time.time()-start_time)
        print(11111)

    def rl_callback(self):
        while True:
            start_time = time.time()
            with torch.no_grad():
                depth_latent = self.obs_class.depth_latent_and_yaw[:, :-2]  # gpu
                yaw = (self.obs_class.depth_latent_and_yaw[:, -2:] * 1.5)
                omega_tensor = torch.tensor(self.omega, dtype=torch.float32, device=device).view(1, -1)
                acc_tensor = torch.tensor(self.acc, dtype=torch.float32, device=device).view(1, -1)

                q_tensor = torch.tensor(self.q, dtype=torch.float32, device=device).view(1, -1)
                qd_tensor = torch.tensor(self.qd, dtype=torch.float32, device=device).view(1, -1)
                contact_tensor = torch.tensor(self.contact, dtype=torch.float32, device=device).view(1, -1)
                # 在这里首先更新下obs的前53个数据
                self.obs_class.insert_proprio(omega_tensor, acc_tensor, yaw, q_tensor, qd_tensor, contact_tensor)
                priprio_last = self.obs_class.obs_buf[:, :args_use.n_proprio]   # gpu
                obs_priv = self.model_use.estimator_model(priprio_last) # gpu
                # print(f"time:{time.time()-start_time}")
                self.obs_class.obs_buf[:, args_use.n_proprio + args_use.n_scan: args_use.n_proprio
                    + args_use.n_scan + args_use.n_priv_explicit] = obs_priv
                actions = self.model_use.actor_model(self.obs_class.obs_buf, True, False, depth_latent)
                # print(f"time:{time.time()-start_time}")
                
                # 发送电机
                q_tensor = torch.tensor(self.q, dtype=torch.float32, device=device).view(1, -1)
                qd_tensor = torch.tensor(self.qd, dtype=torch.float32, device=device).view(1, -1)
                # print(compute_torques(actions, q_tensor, qd_tensor))

                # 下面还需要将tensor类型的结果转化为数组
                # self.motor_msg.tau_des = self.obs_class.compute_torques(actions, self.q, self.qd)

                self.obs_class.insert(actions, priprio_last)
                # 先手动延时0.02s
                # time.sleep(0.02)
            # print(f"time:{time.time()-start_time}")
            # print(11111)



def main(args=None):
    sendZeroFrames()
    rclpy.init(args=args)
    obs_subscriber = ObsSubscriber()
    try:
        obs_subscriber.run()
        Ctrl.quit()
        sys.exit()
    except KeyboardInterrupt:
        pass


parser = argparse.ArgumentParser()
parser.add_argument("--n_obs", type=int, default=753)
parser.add_argument("--n_proprio", type=int, default=53)
parser.add_argument("--n_scan", type=int, default=132)
parser.add_argument("--n_actions", type=int, default=12)
parser.add_argument("--n_priv_latent", type=int, default=29)
parser.add_argument("--n_priv_explicit", type=int, default=9)
parser.add_argument("--n_hist", type=int, default=10)
args_use = parser.parse_args()
if __name__ == "__main__":
    main()
